
local imageio = require "imageio"
local png = require "bindings.libpng"
local alienex = require 'alienex'
local alien = require 'alien'
local cowrap, yield
    = coroutine.wrap, coroutine.yield
local debug = require "debug"
local assert, select
    = assert, select

local buffer_meta = debug.getmetatable(alien.buffer(0))
local function extract_file_pointer(f)
  local mt = debug.getmetatable(f)
  debug.setmetatable(f, buffer_meta)
  local result = f:get(1, 'pointer')
  debug.setmetatable(f, mt)
  return result
end

local function png_error_fn(ptr, msg)
  print("ERROR " .. (msg or '<nil>'))
  error("png error")
end
png_error_fn = alien.callback(png_error_fn, png.error_types)

local function png_warn_fn(ptr, msg)
  print("WARN " .. msg)
end
png_warn_fn = alien.callback(png_warn_fn, png.error_types)

imageio.add_handler{
  pattern = '\137PNG\r\n\026\n';
  reader = function(f)
    local boxed_ptr = alien.buffer(alien.sizeof('pointer'))
    
    local png_ptr = assert(png.create_read_struct(png.LIBPNG_VER_STRING, nil, png_error_fn, png_warn_fn), 'unable to create png read struct')
    
    local info_ptr = assert(png.create_info_struct(png_ptr), 'unable to create png info struct')
    local info = png.info:new(info_ptr)
    
    local reader = alien.callback(
      function(png_ptr, data, size)
        local fetched = f:read(size)
        if fetched and (#fetched == size) then
          alien.memcpy(data, fetched)
        end
      end,
      png.rw_types)
    
    png.set_read_fn(png_ptr, nil, reader)
    
    f:read(8) -- skip magic characters (should already be checked)
    
    local fp = extract_file_pointer(f)
    png.set_sig_bytes(png_ptr, 8)
    png.read_info(png_ptr, info_ptr)
    local fmt, palette
    if (info.color_type == png.COLOR_TYPE_RGB_ALPHA) then
      fmt = 'r8g8b8a8'
    elseif (info.color_type == png.COLOR_TYPE_RGB) then
      fmt = 'r8g8b8'
    elseif (info.color_type == png.COLOR_TYPE_PALETTE) then
      fmt = 'p8'
    elseif (info.color_type == png.COLOR_TYPE_GRAY) then
      fmt = 'p8'
    elseif (info.color_type == png.COLOR_TYPE_GRAY_ALPHA) then
      fmt = 'p8a8'
    end
    
    png.read_update_info(png_ptr, info_ptr)
    
    local row_buf = alien.buffer(info.rowbytes)
    return imageio.create_reader{
      bpp = info.bit_depth;
      format = fmt;
      palette = palette;
      reader = reader;
      
      width = info.width;
      height = info.height;
      pitch = info.rowbytes;
      
      getdata = cowrap(function(self)
        for i = 1, info.height do
          png.read_row(png_ptr, row_buf, nil)
          local rowdata = row_buf:tostring(info.rowbytes)
          yield(i, 1, rowdata)
        end
        return nil
      end);
      
      close = function(self)
      end;
      
      -- png specific
      info = info;
      actual_format = fmt;
      transform_mask = png.TRANSFORM_IDENTITY;
    }
  end;
}

imageio.add_saver{
  name = 'png';
  extensions = {'png'};
  save = function(pixmap, f)
    local png_ptr = assert(png.create_write_struct(png.LIBPNG_VER_STRING, nil, png_error_fn, png_warn_fn), 'unable to create png read struct')
    local info_ptr = assert(png.create_info_struct(png_ptr), 'unable to create png info struct')
    local info = png.info:new(info_ptr)
    
    local format = pixmap.format
    
    local writer = alien.callback(
      function(png_ptr, data, size)
        f:write(alien.buffer(data):tostring(size))
      end,
      png.rw_types)
    
    png.set_write_fn(png_ptr, nil, writer, nil)
    
    local bit_depth, color_type
    
    if (format == 'r8g8b8') then
      bit_depth = 8
      color_type = png.COLOR_TYPE_RGB
    elseif (format == 'r8g8b8a8') then
      bit_depth = 8
      color_type = png.COLOR_TYPE_RGB_ALPHA
    elseif (format == 'b8g8r8') then
      bit_depth = 8
      color_type = png.COLOR_TYPE_RGB
      png.set_bgr(png_ptr)
    elseif (format == 'b8g8r8a8') then
      bit_depth = 8
      color_type = png.COLOR_TYPE_RGB_ALPHA
      png.set_bgr(png_ptr)
    elseif (format == 'b8g8r8x8') then
      bit_depth = 8
      color_type = png.COLOR_TYPE_RGB_ALPHA
      png.set_bgr(png_ptr)
      png.set_strip_alpha(png_ptr)
    elseif (format == 'x8b8g8r8') then
      bit_depth = 8
      color_type = png.COLOR_TYPE_RGB
    else
      error('unsupported pixmap format: ' .. tostring(format))
    end
    
    png.set_IHDR(png_ptr, info_ptr,
      pixmap.width, pixmap.height, bit_depth, color_type,
      png.INTERLACE_NONE, png.COMPRESSION_TYPE_BASE, png.FILTER_TYPE_BASE)
    
    png.write_info(png_ptr, info_ptr)
    
    local row_pointers = pixmap:row_pointers()
    
    if (format == 'x8b8g8r8') then
      png.set_bgr(png_ptr)
      png.set_filler(png_ptr, 1, png.FILLER_BEFORE)
    end
    png.write_image(png_ptr, row_pointers.buffer)
    
    png.write_end(png_ptr, nil)
  end;
}
